import shutil
from pathlib import Path
from typing import Set

def collect_images(source_dir: str, target_dir: str, image_extensions: Set[str]) -> int:
    """
    从源目录递归收集图片文件到目标目录。

    Args:
        source_dir: 要扫描的源目录路径。
        target_dir: 存放收集到的图片的目标目录路径。
        image_extensions: 包含小写图片文件扩展名的集合，例如 {'.jpg', '.png'}。

    Returns:
        成功收集的图片数量。
        
    Raises:
        FileNotFoundError: 如果源目录不存在。
    """
    source_path = Path(source_dir)
    target_path = Path(target_dir)

    if not source_path.is_dir():
        raise FileNotFoundError(f"错误：源目录 '{source_dir}' 不存在或不是一个有效的目录。")

    # 创建目标目录，如果它不存在的话
    target_path.mkdir(parents=True, exist_ok=True)

    collected_count = 0
    print(f"开始扫描目录：'{source_path}'...")

    # 递归遍历源目录中的所有文件
    for file_path in source_path.rglob('*'):
        # 检查文件扩展名是否在支持的图片格式列表中
        if file_path.is_file() and file_path.suffix.lower() in image_extensions:
            try:
                # 构造目标文件路径
                destination_path = target_path / file_path.name
                
                # 为了避免文件名冲突，如果文件已存在，则添加一个唯一的后缀
                counter = 1
                while destination_path.exists():
                    destination_path = target_path / f"{file_path.stem}_{counter}{file_path.suffix}"
                    counter += 1

                # 复制文件并保留元数据
                shutil.copy2(file_path, destination_path)
                print(f"  [+] 成功复制: {file_path.name} -> {destination_path.name}")
                collected_count += 1
            except Exception as e:
                print(f"  [!] 复制文件时出错 {file_path.name}: {e}")

    return collected_count

def main():
    """
    主函数，驱动整个图片收集流程。
    """
    print("--- 一键图片收集器 ---")
    
    # 定义支持的图片文件扩展名
    supported_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
    
    try:
        # 获取用户输入
        source_directory = input("请输入要扫描的源目录路径: ")
        target_directory = input("请输入要将图片存入的目标目录路径 (默认为 'collected_images'): ")

        # 如果用户没有输入目标目录，则使用默认值
        if not target_directory:
            target_directory = "collected_images"

        # 调用核心函数执行收集任务
        count = collect_images(source_directory, target_directory, supported_extensions)

        # 打印最终结果
        if count > 0:
            print(f"\n任务完成！共收集了 {count} 个图片文件。")
            print(f"所有图片已保存至: {Path(target_directory).resolve()}")
        else:
            print("\n扫描完成，但在指定目录中未找到支持的图片文件。")

    except FileNotFoundError as e:
        print(f"\n操作失败: {e}")
    except Exception as e:
        print(f"\n发生了一个未预料的错误: {e}")

if __name__ == "__main__":
    main()